package com.cy.ragbase.service;


import com.cy.ragbase.dto.*;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.model.embedding.EmbeddingModel;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.util.List;
import java.util.stream.Collectors;

@Service
@RequiredArgsConstructor
@Slf4j
public class SearchService {

    private final EmbeddingModel embeddingModel;
    private final MilvusService milvusService;

    /**
     * 执行文档搜索
     */
    public SearchResponse searchDocuments(SearchRequest request) {
        long startTime = System.currentTimeMillis();

        try {
            // 生成查询向量
            Embedding queryEmbedding = embeddingModel.embed(request.getQuery()).content();
            float[] queryVector = queryEmbedding.vector();

            // 在Milvus中搜索相似向量
            List<SearchResult> results = milvusService.searchSimilarChunks(
                    queryVector,
                    request.getTopK(),
                    request.getThreshold()
            );

            long processingTime = System.currentTimeMillis() - startTime;

            // 构建响应
            SearchResponse response = new SearchResponse();
            response.setQuery(request.getQuery());
            response.setResults(results);
            response.setTotalResults(results.size());
            response.setProcessingTimeMs(processingTime);

            log.info("Search completed: query='{}', results={}, time={}ms",
                    request.getQuery(), results.size(), processingTime);

            return response;

        } catch (Exception e) {
            log.error("Search failed for query: {}", request.getQuery(), e);
            throw new RuntimeException("Search failed", e);
        }
    }



    /**
     * 执行混合搜索（结合关键词和语义搜索）
     */
    public SearchResponse hybridSearch(SearchRequest request) {
        // 可以扩展实现混合搜索逻辑
        // 1. 执行关键词搜索
        // 2. 执行语义搜索
        // 3. 合并和重排序结果

        // 目前使用语义搜索
        return searchDocuments(request);
    }
    /**
     * 生成RAG提示词
     */
    public RagPromptResponse generateRagPrompt(RagPromptRequest request) {
        long startTime = System.currentTimeMillis();

        try {
            // 1. 使用现有搜索功能检索相关文档
            SearchRequest searchRequest = new SearchRequest();
            searchRequest.setQuery(request.getQuestion());
            searchRequest.setTopK(request.getTopK());
            searchRequest.setThreshold(request.getThreshold());

            SearchResponse searchResponse = searchDocuments(searchRequest);

            // 2. 转换搜索结果为相关内容
            List<RagPromptResponse.RelevantContent> relevantContents = searchResponse.getResults()
                    .stream()
                    .map(result -> {
                        RagPromptResponse.RelevantContent content = new RagPromptResponse.RelevantContent();
                        content.setDocumentId(result.getDocumentId());
                        content.setContent(result.getContent());
                        content.setScore(result.getScore());
                        content.setChunkIndex(result.getChunkIndex());
                        return content;
                    })
                    .collect(Collectors.toList());

            // 3. 根据提示词类型生成提示词
            String generatedPrompt = buildPrompt(request.getQuestion(), relevantContents, request.getPromptType());

            long processingTime = System.currentTimeMillis() - startTime;

            // 4. 构建响应
            RagPromptResponse response = new RagPromptResponse();
            response.setOriginalQuestion(request.getQuestion());
            response.setGeneratedPrompt(generatedPrompt);
//            response.setRelevantContents(relevantContents);
//            response.setTotalRetrieved(relevantContents.size());
//            response.setProcessingTimeMs(processingTime);

            log.info("RAG prompt generated: question='{}', retrieved={}, time={}ms",
                    request.getQuestion(), relevantContents.size(), processingTime);

            return response;

        } catch (Exception e) {
            log.error("RAG prompt generation failed for question: {}", request.getQuestion(), e);
            throw new RuntimeException("RAG prompt generation failed", e);
        }
    }

    /**
     * 根据提示词类型构建提示词
     */
    private String buildPrompt(String question, List<RagPromptResponse.RelevantContent> relevantContents,
                               RagPromptRequest.PromptType promptType) {
        StringBuilder prompt = new StringBuilder();

        switch (promptType) {
            case CODE_GENERATION:
                prompt.append("# 代码生成任务\n\n");
                prompt.append("## 用户需求\n");
                prompt.append(question).append("\n\n");
                prompt.append("## 相关参考文档\n");
                break;

            case GENERAL_QA:
                prompt.append("# 问答任务\n\n");
                prompt.append("## 问题\n");
                prompt.append(question).append("\n\n");
                prompt.append("## 参考资料\n");
                break;

            case DOCUMENTATION:
                prompt.append("# 文档说明任务\n\n");
                prompt.append("## 说明需求\n");
                prompt.append(question).append("\n\n");
                prompt.append("## 相关文档内容\n");
                break;
        }

        // 添加检索到的相关文档内容
        for (int i = 0; i < relevantContents.size(); i++) {
            RagPromptResponse.RelevantContent content = relevantContents.get(i);
            prompt.append("### 参考文档 ").append(i + 1)
                    .append(" (相似度: ").append(String.format("%.2f", content.getScore())).append(")\n");
            prompt.append("```\n");
            prompt.append(content.getContent());
            prompt.append("\n```\n\n");
        }

        // 添加任务指令
        switch (promptType) {
            case CODE_GENERATION:
                prompt.append("## 任务要求\n");
                prompt.append("请根据上述需求和参考文档，生成相应的代码。要求：\n");
                prompt.append("1. 代码应该清晰、可读性强\n");
                prompt.append("2. 适当添加注释说明\n");
                prompt.append("3. 遵循最佳实践和设计模式\n");
                prompt.append("4. 如果参考文档中有相关实现，请参考其风格和结构\n");
                break;

            case GENERAL_QA:
                prompt.append("## 回答要求\n");
                prompt.append("请根据上述问题和参考资料，提供准确、详细的回答。\n");
                break;

            case DOCUMENTATION:
                prompt.append("## 输出要求\n");
                prompt.append("请根据相关文档内容，生成清晰、完整的说明文档。\n");
                break;
        }

        return prompt.toString();
    }
}